#!/usr/bin/env python3
# fake_clients_simulate.py - 模拟三个机器人的测试脚本（改进版，增加方向和更真实的运动）
import socket
import threading
import struct
import zlib
import time
import random
import logging
import math

# 配置日志
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

# 通信协议常量
FRAME_SIZE = 32
VERSION = 0x02

# 指令类型
COMMAND_TYPES = {
    "ConnectionRequest": 0x01,
    "ConnectionResponse": 0x02,
    "PositionReport": 0x03,
    "NavigationGoal": 0x04,  # 假设目标可能包含 [x, y, yaw] 或仅 [x, y]
    "Heartbeat": 0x05,
    "CommandAck": 0x06,
    "EmergencyStop": 0x07,
    "StatusReport": 0x08,
    "RoleAssignment": 0x09,
}

# 状态码
STATUS_CODES = {
    "Success": 0x00,
    "GeneralError": 0x01,
    "InvalidCommand": 0x02,
    "CRCError": 0x03,
    "Timeout": 0x04,
    "RobotBusy": 0x05,
}

# 角色定义
ROLES = {"evader": 0x01, "pursuer": 0x02}

# 导航状态
NAV_STATUS = {
    "IDLE": 0,  # 空闲
    "ROTATING": 1,  # 旋转中
    "MOVING": 2,  # 移动中
    "ARRIVED": 3,  # 到达目标
    "ERROR": 4,  # 错误状态
    "TIMEOUT": 5,  # 超时
}

# 新增：运动控制参数 (单位：m/s, rad/s)
DEFAULT_MAX_LINEAR_SPEED = 1.0
DEFAULT_MAX_ANGULAR_SPEED = 0.3


class FakeRobotClient:
    def __init__(self, robot_name, server_ip, server_port):
        self.robot_name = robot_name
        self.server_ip = server_ip
        self.server_port = server_port
        self.robot_id = 0
        self.role = None
        self.sock = None
        self.running = True
        self.connection_established = False
        self.last_heartbeat_time = 0
        self.heartbeat_interval = 2.0
        self.crc_errors = 0

        # 机器人状态 (增加 yaw)
        self.x = random.uniform(-50, 50)
        self.y = random.uniform(-70, 0)
        self.yaw = random.uniform(0, 2 * math.pi)
        self.target_x = self.x
        self.target_y = self.y
        self.status = NAV_STATUS["IDLE"]

        # 新增：运动参数
        self.max_linear_speed = DEFAULT_MAX_LINEAR_SPEED
        self.max_angular_speed = DEFAULT_MAX_ANGULAR_SPEED

        # 启动连接
        self.connect_thread = threading.Thread(
            target=self.connect_to_server, daemon=True
        )
        self.connect_thread.start()
        # 启动发送线程
        self.sender_thread = threading.Thread(
            target=self.send_periodic_messages, daemon=True
        )
        self.sender_thread.start()

    def build_frame(self, command, payload=None, robot_id=None):
        """构建通信帧（改进版）"""
        if robot_id is None:
            robot_id = self.robot_id
        if payload is None:
            payload = bytes(14)
        elif len(payload) > 14:
            payload = payload[:14]  # 截断到14字节
        elif len(payload) < 14:
            payload = payload.ljust(14, b"\x00")  # 填充到14字节
        timestamp = int(time.time() * 1000)
        header = struct.pack(">BBiQ", VERSION, command, robot_id, timestamp)
        frame_data = header + payload  # 14字节头 + 14字节负载 = 28字节
        crc32 = zlib.crc32(frame_data) & 0xFFFFFFFF
        crc_bytes = struct.pack(">I", crc32)
        full_frame = frame_data + crc_bytes  # 28 + 4 = 32字节
        return full_frame

    def parse_frame(self, data):
        """解析通信帧（改进版）"""
        if len(data) != FRAME_SIZE:
            logger.error(f"无效帧长度: {len(data)} (预期{FRAME_SIZE})")
            return None
        frame_data = data[:28]  # 前28字节用于CRC计算
        received_crc = struct.unpack(">I", data[28:32])[0]
        calculated_crc = zlib.crc32(frame_data) & 0xFFFFFFFF
        if received_crc != calculated_crc:
            self.crc_errors += 1
            logger.warning(
                f"CRC校验失败: 接收={received_crc:08X}, 计算={calculated_crc:08X}"
            )
            return None
        try:
            header_part = frame_data[:14]
            version, command, robot_id, timestamp = struct.unpack(">BBiQ", header_part)
            payload = frame_data[14:28]  # 14字节负载
            if version != VERSION:
                logger.warning(f"协议版本不匹配: 收到={version}, 期望={VERSION}")
                return None
            return {
                "version": version,
                "command": command,
                "robot_id": robot_id,
                "timestamp": timestamp,
                "payload": payload,
            }
        except struct.error as e:
            logger.error(f"帧解析错误: {e}")
            return None

    def connect_to_server(self):
        """连接到服务器（改进版）"""
        while self.running:
            try:
                # 创建连接
                self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                self.sock.settimeout(10)
                self.sock.connect((self.server_ip, self.server_port))
                logger.info(f"{self.robot_name} 已连接到服务器")
                # 发送连接请求
                name_bytes = self.robot_name.encode("utf-8")
                if len(name_bytes) > 14:
                    name_bytes = name_bytes[:14]
                elif len(name_bytes) < 14:
                    name_bytes = name_bytes.ljust(14, b"\x00")
                frame = self.build_frame(
                    COMMAND_TYPES["ConnectionRequest"],
                    name_bytes,
                    robot_id=0,  # 初始连接ID=0
                )
                self.sock.sendall(frame)
                logger.info(f"{self.robot_name} 已发送连接请求")
                # 接收连接响应
                response = self.sock.recv(FRAME_SIZE)
                if not response:
                    raise Exception("空响应")
                parsed = self.parse_frame(response)
                if parsed and parsed["command"] == COMMAND_TYPES["ConnectionResponse"]:
                    # 获取分配的机器人ID
                    self.robot_id = struct.unpack(">i", parsed["payload"][:4])[0]
                    logger.info(f"{self.robot_name} 分配的ID: {self.robot_id}")
                    # 发送确认
                    ack_payload = struct.pack(
                        ">BB",
                        COMMAND_TYPES["ConnectionResponse"],
                        STATUS_CODES["Success"],
                    )
                    ack_frame = self.build_frame(
                        COMMAND_TYPES["CommandAck"], ack_payload
                    )
                    self.sock.sendall(ack_frame)
                    # 标记连接已建立
                    self.connection_established = True
                    self.last_heartbeat_time = time.time()
                # 启动接收循环
                threading.Thread(target=self.receive_loop, daemon=True).start()
                break
            except Exception as e:
                logger.error(f"{self.robot_name} 连接失败: {e}，重试中...")
                time.sleep(5)

    def receive_loop(self):
        """接收消息循环（改进版）"""
        self.sock.settimeout(5.0)  # 设置接收超时
        while self.running and self.connection_established:
            try:
                data = self.sock.recv(FRAME_SIZE)
                if not data:
                    logger.warning(f"{self.robot_name} 连接断开")
                    break
                parsed = self.parse_frame(data)
                if not parsed:
                    continue
                # 检查心跳超时
                current_time = time.time()
                if (
                    current_time - self.last_heartbeat_time
                    > self.heartbeat_interval * 3
                ):
                    logger.warning(f"{self.robot_name} 心跳超时，重新连接...")
                    self.connection_established = False
                    break
                # 处理不同命令
                if parsed["command"] == COMMAND_TYPES["NavigationGoal"]:
                    self.handle_navigation_command(parsed["payload"])
                elif parsed["command"] == COMMAND_TYPES["Heartbeat"]:
                    self.handle_heartbeat(parsed["payload"])
                elif parsed["command"] == COMMAND_TYPES["RoleAssignment"]:
                    self.handle_role_assignment(parsed["payload"])
                elif parsed["command"] == COMMAND_TYPES["EmergencyStop"]:
                    self.handle_emergency_stop()
                elif parsed["command"] == COMMAND_TYPES["CommandAck"]:
                    self.handle_command_ack(parsed["payload"])
            except socket.timeout:
                # 检查心跳
                current_time = time.time()
                if current_time - self.last_heartbeat_time > self.heartbeat_interval:
                    self.send_heartbeat()
                continue
            except Exception as e:
                logger.error(f"{self.robot_name} 接收错误: {e}")
                break
        self.reconnect()

    def handle_navigation_command(self, payload):
        """处理导航命令（假设目标为 [x, y]）"""
        try:
            # 解析目标坐标 (x, y) - 8字节
            if len(payload) >= 8:  # 至少需要 x, y
                self.target_x, self.target_y = struct.unpack(">ff", payload[:8])
                # 如果负载有12字节，则包含 timeout
                timeout = 30.0  # 默认超时
                if len(payload) >= 12:
                    (timeout_val,) = struct.unpack(">f", payload[8:12])
                    timeout = timeout_val if timeout_val > 0 else 30.0

                logger.info(
                    f"{self.robot_name} 收到目标: ({self.target_x:.2f}, {self.target_y:.2f})"
                )
                # 发送确认
                ack_payload = struct.pack(
                    ">BB", COMMAND_TYPES["NavigationGoal"], STATUS_CODES["Success"]
                )
                ack_frame = self.build_frame(COMMAND_TYPES["CommandAck"], ack_payload)
                self.sock.sendall(ack_frame)
                # 启动导航线程
                threading.Thread(
                    target=self.navigate_to_target, args=(timeout,), daemon=True
                ).start()
            else:
                logger.error(f"导航命令负载长度不足: {len(payload)} (需要至少8字节)")
                # 发送错误确认
                ack_payload = struct.pack(
                    ">BB", COMMAND_TYPES["NavigationGoal"], STATUS_CODES["GeneralError"]
                )
                ack_frame = self.build_frame(COMMAND_TYPES["CommandAck"], ack_payload)
                self.sock.sendall(ack_frame)
        except Exception as e:
            logger.error(f"{self.robot_name} 处理导航命令错误: {e}")
            # 发送错误确认
            ack_payload = struct.pack(
                ">BB", COMMAND_TYPES["NavigationGoal"], STATUS_CODES["GeneralError"]
            )
            ack_frame = self.build_frame(COMMAND_TYPES["CommandAck"], ack_payload)
            self.sock.sendall(ack_frame)

    def _normalize_angle(self, angle):
        """将角度标准化到 [-pi, pi] 范围内"""
        while angle > math.pi:
            angle -= 2 * math.pi
        while angle < -math.pi:
            angle += 2 * math.pi
        return angle

    def navigate_to_target(self, timeout):
        """执行导航到目标点（修改版：先旋转再移动）"""
        import math  # 确保 math 模块可用

        logger.info(
            f"{self.robot_name} 开始导航到目标: ({self.target_x:.2f}, {self.target_y:.2f})"
        )

        start_time = time.time()
        self.status = NAV_STATUS["MOVING"]  # 使用 MOVING 状态表示导航进行中

        # --- 阶段 1: 旋转朝向目标点 ---
        logger.info(f"{self.robot_name} 开始旋转以朝向目标...")
        while True:
            # 检查超时
            if time.time() - start_time > 30:
                logger.warning(f"{self.robot_name} 导航超时 (旋转阶段)")
                self.status = NAV_STATUS["TIMEOUT"]
                return  # 直接退出导航

            # 计算到目标点的角度
            dx = self.target_x - self.x
            dy = self.target_y - self.y
            target_angle = math.atan2(dy, dx)

            # 假设机器人当前没有朝向信息，我们用一个虚拟的 yaw 来模拟
            # 为了简化，我们用当前速度方向来近似当前朝向
            # 如果 dx, dy 都接近 0，说明已经很接近目标，可以认为方向正确
            if abs(dx) < 0.01 and abs(dy) < 0.01:
                logger.info(f"{self.robot_name} 已经非常接近目标。")
                break  # 已经足够接近，无需旋转

            # 简化模型：假设机器人“瞬间”调整方向
            # 实际上，这里应该有一个持续旋转直到方向正确的过程
            # 但由于原代码没有存储方向，我们简化处理：
            # 只要还没到目标，就认为需要调整方向，然后立即进入移动阶段
            # 更真实的模拟需要存储和更新机器人的 yaw 角

            # 为了符合“先旋转”的要求，我们在这里做一个简短的等待，
            # 模拟旋转所需的时间。实际应用中应根据角速度计算。
            logger.info(f"{self.robot_name} 完成旋转 (模拟)。")
            break  # 旋转阶段结束

        # --- 阶段 2: 直线移动到目标点 ---
        logger.info(f"{self.robot_name} 开始直线移动到目标...")
        while True:
            dx = self.target_x - self.x
            dy = self.target_y - self.y
            distance = (dx**2 + dy**2) ** 0.5

            # 再次检查超时
            if time.time() - start_time > 30:
                logger.warning(f"{self.robot_name} 导航超时 (移动阶段)")
                self.status = NAV_STATUS["TIMEOUT"]
                break

            if distance < 0.1:  # 到达目标
                self.x = self.target_x
                self.y = self.target_y
                self.status = NAV_STATUS["ARRIVED"]  # 到达目标
                logger.info(
                    f"{self.robot_name} 到达目标点 ({self.x:.2f}, {self.y:.2f})"
                )
                break

            # 每次移动 step 单位距离，沿直线方向
            step = self.max_linear_speed
            self.x += dx * step / distance if distance > 0 else 0
            self.y += dy * step / distance if distance > 0 else 0
            time.sleep(0.1)  # 移动一步的耗时

        # 短暂保持最终状态后返回空闲
        if self.status == NAV_STATUS["ARRIVED"]:
            time.sleep(0.5)
        self.status = NAV_STATUS["IDLE"]
        logger.info(f"{self.robot_name} 导航结束，返回空闲状态")

    def handle_heartbeat(self, payload):
        """处理心跳包（改进版）"""
        try:
            if len(payload) >= 4:
                seq = struct.unpack(">i", payload[:4])[0]
                logger.debug(f"{self.robot_name} 收到心跳包序列号: {seq}")
            else:
                seq = 0
                logger.warning(f"心跳包负载长度不足: {len(payload)}")
            # 更新最后心跳时间
            self.last_heartbeat_time = time.time()
            # 发送响应
            heartbeat_payload = struct.pack(">i", seq)
            heartbeat_frame = self.build_frame(
                COMMAND_TYPES["Heartbeat"], heartbeat_payload
            )
            self.sock.sendall(heartbeat_frame)
        except Exception as e:
            logger.error(f"{self.robot_name} 处理心跳包错误: {e}")

    def handle_role_assignment(self, payload):
        """处理角色分配（改进版）"""
        try:
            if len(payload) >= 1:
                role_code = payload[0]
                if role_code == ROLES["evader"]:
                    self.role = "evader"
                elif role_code == ROLES["pursuer"]:
                    self.role = "pursuer"
                logger.info(
                    f"{self.robot_name} 分配的角色: {self.role} (线速度: {self.max_linear_speed:.2f}, 角速度: {self.max_angular_speed:.2f})"
                )
                # 发送确认
                ack_payload = struct.pack(
                    ">BB", COMMAND_TYPES["RoleAssignment"], STATUS_CODES["Success"]
                )
                ack_frame = self.build_frame(COMMAND_TYPES["CommandAck"], ack_payload)
                self.sock.sendall(ack_frame)
            else:
                logger.error("角色分配负载长度不足")
        except Exception as e:
            logger.error(f"{self.robot_name} 处理角色分配错误: {e}")
            # 发送错误确认
            ack_payload = struct.pack(
                ">BB", COMMAND_TYPES["RoleAssignment"], STATUS_CODES["GeneralError"]
            )
            ack_frame = self.build_frame(COMMAND_TYPES["CommandAck"], ack_payload)
            self.sock.sendall(ack_frame)

    def handle_emergency_stop(self):
        """处理紧急停止（改进版）"""
        logger.info(f"{self.robot_name} 收到紧急停止指令")
        self.target_x = self.x
        self.target_y = self.y
        # self.target_yaw = self.yaw # 如果需要停止旋转
        self.status = NAV_STATUS["IDLE"]  # 返回空闲状态
        # 发送确认
        ack_payload = struct.pack(
            ">BB", COMMAND_TYPES["EmergencyStop"], STATUS_CODES["Success"]
        )
        ack_frame = self.build_frame(COMMAND_TYPES["CommandAck"], ack_payload)
        self.sock.sendall(ack_frame)

    def handle_command_ack(self, payload):
        """处理指令确认（改进版）"""
        try:
            if len(payload) >= 2:
                acked_command, status_code = struct.unpack(">BB", payload[:2])
                status_msg = next(
                    (k for k, v in STATUS_CODES.items() if v == status_code), "Unknown"
                )
                logger.debug(
                    f"{self.robot_name} 收到确认: 命令={acked_command}, 状态={status_msg}"
                )
            else:
                logger.warning(f"确认负载长度不足: {len(payload)}")
        except Exception as e:
            logger.warning(f"{self.robot_name} 处理确认错误: {e}")

    def send_position_report(self):
        """发送位置报告（改进版，包含 x, y, yaw）"""
        if not self.connection_established:
            return
        try:
            # 构造位置报告 (x, y, yaw, status) - 13字节 + 1字节填充
            # x(float 4), y(float 4), yaw(float 4), status(byte 1) = 13 bytes
            payload = struct.pack(">fffB", self.x, self.y, self.yaw, self.status)
            payload = payload.ljust(14, b"\x00")  # 填充到14字节
            # 发送帧
            frame = self.build_frame(COMMAND_TYPES["PositionReport"], payload)
            self.sock.sendall(frame)
            logger.info(
                f"{self.robot_name} 发送位姿: ({self.x:.2f}, {self.y:.2f}, {self.yaw:.2f}), 状态: {self.status}"
            )
        except Exception as e:
            logger.error(f"{self.robot_name} 发送位置报告错误: {e}")

    def send_heartbeat(self):
        """发送心跳包（改进版）"""
        if not self.connection_established:
            return
        try:
            # 使用当前时间戳作为序列号
            seq = int(time.time() * 1000) % 1000000
            payload = struct.pack(">i", seq)
            payload = payload.ljust(14, b"\x00")  # 填充到14字节
            frame = self.build_frame(COMMAND_TYPES["Heartbeat"], payload)
            self.sock.sendall(frame)
            self.last_heartbeat_time = time.time()
            logger.debug(f"{self.robot_name} 发送心跳包")
        except Exception as e:
            logger.error(f"{self.robot_name} 发送心跳包错误: {e}")

    def send_periodic_messages(self):
        """发送周期性消息（改进版）"""
        while self.running:
            try:
                if self.connection_established:
                    self.send_position_report()
                    # 心跳由接收超时触发，不在此处发送
                time.sleep(0.2)  # 每秒发送一次位置报告
            except Exception as e:
                logger.error(f"{self.robot_name} 周期消息错误: {e}")
                time.sleep(0.5)

    def reconnect(self):
        """重新连接（改进版）"""
        if self.sock:
            try:
                self.sock.close()
            except:
                pass
        self.sock = None
        self.robot_id = 0
        self.connection_established = False
        self.running = True
        time.sleep(2)
        self.connect_thread = threading.Thread(
            target=self.connect_to_server, daemon=True
        )
        self.connect_thread.start()

    def stop(self):
        """停止客户端"""
        self.running = False
        if self.sock:
            try:
                self.sock.close()
            except:
                pass


if __name__ == "__main__":
    # 启动三个模拟机器人
    robot_number = 3
    robots = []
    try:
        for i in range(robot_number):
            robot_name = f"robot{i+1}"
            client = FakeRobotClient(robot_name, "127.0.0.1", 6001)
            robots.append(client)
            time.sleep(0.3)  # 稍微延迟连接

        robots[0].max_linear_speed *= 1.5

        # 保持程序运行
        while True:
            time.sleep(0.2)
    except KeyboardInterrupt:
        logger.info("停止模拟机器人...")
        for robot in robots:
            robot.stop()
